'''
H5 Event Representation Visualization
读取main.py生成的H5格式事件表示数据并可视化
支持StackedHistogram和MixedDensityStack格式
'''
import h5py
import cv2
import numpy as np
import os
import glob
import argparse
from tqdm import tqdm
from pathlib import Path
import matplotlib.pyplot as plt

np.random.seed(0)

def parse_argument():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    
    ## 目录参数
    parser.add_argument('--h5_file_path', type=str, required=True,
                        help="J:\dataset\\nuclear_power\\nuclear_H5_process_new\\val\\recording_2024-06-05_10-44-23\event_representations_v2\stacked_histogram_dt=50_nbins=10/event_representations.h5" )
    parser.add_argument('--save_dir', type=str,
                        default='./h5_visualizations/',
                        help='可视化图像输出目录')
    
    ## 可视化参数
    parser.add_argument('--vis_mode', type=str, default='stacked',
                        choices=['stacked', 'channels', 'mixed_density', 'summary'],
                        help='可视化模式')
    parser.add_argument('--color_mode', type=str, default='gray',
                        choices=['gray', 'jet', 'hot', 'viridis', 'plasma'],
                        help='颜色映射模式')
    
    ## 处理参数
    parser.add_argument('--max_frames', type=int, default=50,
                        help='最大处理帧数')
    parser.add_argument('--save_images', action='store_true',
                        help='保存图像而不是实时显示')
    parser.add_argument('--show_info', action='store_true',
                        help='显示详细信息')
    
    return parser

def filter_outliers(image, percentile=7):
    """滤波异常值"""
    if len(image.shape) == 3:
        # 如果是多通道图像，处理每个通道
        filtered_channels = []
        for c in range(image.shape[0]):
            channel = image[c]
            low_val = np.percentile(channel, percentile)
            high_val = np.percentile(channel, 100 - percentile)
            filtered_channel = np.clip(channel, low_val, high_val)
            if high_val > low_val:
                filtered_channel = (filtered_channel - low_val) / (high_val - low_val)
            filtered_channels.append(filtered_channel)
        return np.stack(filtered_channels, axis=0)
    else:
        # 单通道处理
        low_val = np.percentile(image, percentile)
        high_val = np.percentile(image, 100 - percentile)
        filtered_image = np.clip(image, low_val, high_val)
        if high_val > low_val:
            filtered_image = (filtered_image - low_val) / (high_val - low_val)
        return filtered_image

def create_viridis_colormap():
    """创建Viridis颜色映射"""
    viridis_colors = np.array([
        [68, 1, 84], [59, 82, 139], [33, 144, 140], 
        [93, 201, 99], [253, 231, 37]
    ]) / 255.0
    
    lut = np.zeros((256, 3))
    for i in range(256):
        value = i / 255.0
        color_idx = value * (len(viridis_colors) - 1)
        idx_low = int(np.floor(color_idx))
        idx_high = min(idx_low + 1, len(viridis_colors) - 1)
        alpha = color_idx - idx_low
        lut[i] = (1 - alpha) * viridis_colors[idx_low] + alpha * viridis_colors[idx_high]
    
    return lut

class EventRepresentationVisualizer:
    """事件表示可视化器"""
    
    def __init__(self, color_mode='viridis'):
        self.color_mode = color_mode
        
        # 预计算颜色映射
        if color_mode == 'viridis':
            self.viridis_lut = create_viridis_colormap()
    
    def apply_colormap(self, image):
        """应用颜色映射"""
        if self.color_mode == 'gray':
            return (image * 255).astype(np.uint8)
        
        # 转换为uint8
        image_uint8 = (image * 255).astype(np.uint8)
        
        if self.color_mode == 'viridis':
            # 使用预计算的Viridis查找表
            colored = self.viridis_lut[image_uint8]
            return (colored * 255)[:, :, ::-1].astype(np.uint8)  # RGB to BGR
        elif self.color_mode == 'jet':
            return cv2.applyColorMap(image_uint8, cv2.COLORMAP_JET)
        elif self.color_mode == 'hot':
            return cv2.applyColorMap(image_uint8, cv2.COLORMAP_HOT)
        elif self.color_mode == 'plasma':
            return cv2.applyColorMap(image_uint8, cv2.COLORMAP_PLASMA)
        else:
            return cv2.applyColorMap(image_uint8, cv2.COLORMAP_JET)
    
    def visualize_stacked_histogram(self, frame_data):
        """可视化堆叠直方图 - 显示所有通道的组合"""
        if len(frame_data.shape) == 3:  # (C, H, W)
            # 对所有通道求和或平均
            combined = np.mean(frame_data, axis=0)
        else:  # (H, W)
            combined = frame_data
        
        # 滤波和归一化
        filtered = filter_outliers(combined, 7)
        
        return self.apply_colormap(filtered)
    
    def visualize_channels_separate(self, frame_data):
        """分别可视化每个通道"""
        if len(frame_data.shape) != 3:
            # 如果不是多通道，返回单个图像
            filtered = filter_outliers(frame_data, 7)
            return [self.apply_colormap(filtered)]
        
        channel_images = []
        for c in range(frame_data.shape[0]):
            channel = frame_data[c]
            filtered = filter_outliers(channel, 7)
            colored = self.apply_colormap(filtered)
            channel_images.append(colored)
        
        return channel_images
    
    def visualize_mixed_density(self, frame_data):
        """可视化混合密度堆叠"""
        if len(frame_data.shape) == 3:  # (C, H, W)
            # 创建RGB图像，使用不同通道
            if frame_data.shape[0] >= 3:
                # 使用前3个通道作为RGB
                rgb_image = np.zeros((frame_data.shape[1], frame_data.shape[2], 3))
                for i in range(3):
                    channel = filter_outliers(frame_data[i], 7)
                    rgb_image[:, :, i] = channel
                
                return (rgb_image * 255).astype(np.uint8)
            else:
                # 通道数不足，使用灰度
                combined = np.mean(frame_data, axis=0)
                filtered = filter_outliers(combined, 7)
                return self.apply_colormap(filtered)
        else:
            filtered = filter_outliers(frame_data, 7)
            return self.apply_colormap(filtered)
    
    def create_summary_view(self, frame_data):
        """创建摘要视图 - 显示多个可视化结果"""
        if len(frame_data.shape) != 3:
            # 单通道数据
            filtered = filter_outliers(frame_data, 7)
            return self.apply_colormap(filtered)
        
        num_channels = frame_data.shape[0]
        
        if num_channels <= 4:
            # 创建2x2网格
            grid_size = 2
            grid_image = np.zeros((frame_data.shape[1] * grid_size, 
                                 frame_data.shape[2] * grid_size, 3), dtype=np.uint8)
            
            # 填充网格
            positions = [(0, 0), (0, 1), (1, 0), (1, 1)]
            for i in range(min(num_channels, 4)):
                row, col = positions[i]
                channel = filter_outliers(frame_data[i], 7)
                colored = self.apply_colormap(channel)
                
                if len(colored.shape) == 2:
                    colored = cv2.cvtColor(colored, cv2.COLOR_GRAY2BGR)
                
                start_row = row * frame_data.shape[1]
                end_row = (row + 1) * frame_data.shape[1]
                start_col = col * frame_data.shape[2]
                end_col = (col + 1) * frame_data.shape[2]
                
                grid_image[start_row:end_row, start_col:end_col] = colored
            
            return grid_image
        else:
            # 通道太多，显示组合视图
            return self.visualize_stacked_histogram(frame_data)

def analyze_h5_structure(h5_file_path):
    """分析H5文件结构"""
    print(f'\n=== 分析H5文件结构 ===')
    print(f'文件: {h5_file_path}')
    print(f'文件大小: {os.path.getsize(h5_file_path)/1e6:.3f}MB')
    
    with h5py.File(h5_file_path, 'r') as f:
        print(f'根级键: {list(f.keys())}')
        
        # 查找数据集
        data_key = 'data' if 'data' in f else list(f.keys())[0]
        dataset = f[data_key]
        
        print(f'数据集键: {data_key}')
        print(f'数据形状: {dataset.shape}')
        print(f'数据类型: {dataset.dtype}')
        
        # 显示属性
        if hasattr(dataset, 'attrs') and len(dataset.attrs) > 0:
            print(f'\n=== 数据集属性 ===')
            for key in dataset.attrs:
                print(f'  {key}: {dataset.attrs[key]}')
        
        # 分析数据特征
        print(f'\n=== 数据特征分析 ===')
        if len(dataset.shape) == 4:  # (N, C, H, W)
            print(f'总帧数: {dataset.shape[0]}')
            print(f'通道数: {dataset.shape[1]}')
            print(f'图像尺寸: {dataset.shape[3]}x{dataset.shape[2]}')
            
            # 采样分析前几帧
            sample_frame = dataset[0]
            print(f'样本帧范围: [{np.min(sample_frame):.3f}, {np.max(sample_frame):.3f}]')
            print(f'样本帧均值: {np.mean(sample_frame):.3f}')
            print(f'非零像素比例: {np.count_nonzero(sample_frame) / sample_frame.size:.3f}')
            
        return data_key, dataset.shape, dataset.dtype

def visualize_h5_event_representations(h5_file_path, save_dir="./h5_visualizations/", 
                                     vis_mode='stacked', color_mode='viridis',
                                     max_frames=50, save_images=False, show_info=True):
    """可视化H5事件表示文件"""
    
    if not os.path.exists(h5_file_path):
        print(f"错误: 找不到文件 {h5_file_path}")
        return
    
    # 分析文件结构
    if show_info:
        data_key, data_shape, data_dtype = analyze_h5_structure(h5_file_path)
    
    # 创建输出目录
    if save_images:
        base_name = Path(h5_file_path).stem
        output_dir = os.path.join(save_dir, base_name)
        os.makedirs(output_dir, exist_ok=True)
        print(f'图像将保存到: {output_dir}')
    
    # 创建可视化器
    visualizer = EventRepresentationVisualizer(color_mode=color_mode)
    
    try:
        with h5py.File(h5_file_path, 'r') as f:
            # 获取数据集
            data_key = 'data' if 'data' in f else list(f.keys())[0]
            dataset = f[data_key]
            
            # 确定处理帧数
            total_frames = dataset.shape[0]
            frames_to_process = min(total_frames, max_frames)
            
            print(f"\n=== 开始可视化 ===")
            print(f"可视化模式: {vis_mode}")
            print(f"颜色映射: {color_mode}")
            print(f"总帧数: {total_frames}")
            print(f"处理帧数: {frames_to_process}")
            
            # 设置matplotlib (如果不保存图像)
            if not save_images:
                plt.ion()
                if vis_mode == 'channels':
                    # 为多通道显示创建子图
                    num_channels = dataset.shape[1] if len(dataset.shape) == 4 else 1
                    cols = min(num_channels, 4)
                    rows = (num_channels + cols - 1) // cols
                    fig, axes = plt.subplots(rows, cols, figsize=(cols*4, rows*3))
                    if num_channels == 1:
                        axes = [axes]
                    elif rows == 1:
                        axes = [axes]
                    else:
                        axes = axes.flatten()
                else:
                    fig, ax = plt.subplots(figsize=(12, 8))
            
            # 处理每一帧
            for i in tqdm(range(frames_to_process), desc="处理帧"):
                try:
                    # 读取帧数据
                    frame_data = dataset[i]
                    
                    # 根据可视化模式处理
                    if vis_mode == 'stacked':
                        vis_image = visualizer.visualize_stacked_histogram(frame_data)
                        vis_images = [vis_image]
                        titles = [f'Frame {i} - Stacked View']
                        
                    elif vis_mode == 'channels':
                        vis_images = visualizer.visualize_channels_separate(frame_data)
                        titles = [f'Frame {i} - Channel {c}' for c in range(len(vis_images))]
                        
                    elif vis_mode == 'mixed_density':
                        vis_image = visualizer.visualize_mixed_density(frame_data)
                        vis_images = [vis_image]
                        titles = [f'Frame {i} - Mixed Density']
                        
                    elif vis_mode == 'summary':
                        vis_image = visualizer.create_summary_view(frame_data)
                        vis_images = [vis_image]
                        titles = [f'Frame {i} - Summary View']
                    
                    # 保存或显示
                    if save_images:
                        for idx, (vis_image, title) in enumerate(zip(vis_images, titles)):
                            if vis_mode == 'channels':
                                filename = f'frame_{i:06d}_channel_{idx:02d}.png'
                            else:
                                filename = f'frame_{i:06d}_{vis_mode}.png'
                            filepath = os.path.join(output_dir, filename)
                            cv2.imwrite(filepath, vis_image)
                    else:
                        # 实时显示
                        if vis_mode == 'channels':
                            for idx, (vis_image, title) in enumerate(zip(vis_images, titles)):
                                if idx < len(axes):
                                    axes[idx].clear()
                                    if len(vis_image.shape) == 3:
                                        axes[idx].imshow(cv2.cvtColor(vis_image, cv2.COLOR_BGR2RGB))
                                    else:
                                        axes[idx].imshow(vis_image, cmap='gray')
                                    axes[idx].set_title(title)
                                    axes[idx].axis('off')
                        else:
                            ax.clear()
                            vis_image = vis_images[0]
                            if len(vis_image.shape) == 3:
                                ax.imshow(cv2.cvtColor(vis_image, cv2.COLOR_BGR2RGB))
                            else:
                                ax.imshow(vis_image, cmap='gray')
                            ax.set_title(titles[0])
                            ax.axis('off')
                        
                        plt.tight_layout()
                        plt.pause(0.1)
                        
                        # 检查窗口是否关闭
                        if not plt.get_fignums():
                            break
                
                except Exception as e:
                    print(f"处理第 {i} 帧时出错: {e}")
                    continue
            
            if not save_images:
                plt.ioff()
                plt.show()
            
            print(f"\n完成！处理了 {frames_to_process} 帧")
            if save_images:
                print(f"图像已保存到: {output_dir}")
                
    except Exception as e:
        print(f"读取H5文件时出错: {e}")

def main():
    """主函数"""
    args, _ = parse_argument().parse_known_args()
    
    print(f"参数配置:")
    print(f"  H5文件: {args.h5_file_path}")
    print(f"  保存目录: {args.save_dir}")
    print(f"  可视化模式: {args.vis_mode}")
    print(f"  颜色映射: {args.color_mode}")
    print(f"  最大帧数: {args.max_frames}")
    print(f"  保存图像: {args.save_images}")
    
    # 可视化H5文件
    visualize_h5_event_representations(
        h5_file_path=args.h5_file_path,
        save_dir=args.save_dir,
        vis_mode=args.vis_mode,
        color_mode=args.color_mode,
        max_frames=args.max_frames,
        save_images=args.save_images,
        show_info=args.show_info
    )

if __name__ == '__main__':
    # 如果直接运行，可以修改这里的路径
    h5_file_path = "J:\dataset\\nuclear_power\\nuclear_H5_process_new\\val\\recording_2024-06-05_10-44-23\event_representations_v2\stacked_histogram_dt=50_nbins=10/event_representations.h5"  # 替换为你的H5文件路径
    
    # 检查是否有命令行参数
    import sys
    if len(sys.argv) > 1:
        main()
    else:
        # 直接运行模式
        print("直接运行模式")
        print("请修改 h5_file_path 变量为你的H5文件路径")
        print("或使用命令行模式:")
        print("python h5_event_repr_visualizer.py --h5_file_path your_file.h5")
        
        # 示例：如果文件存在，自动可视化
        if os.path.exists(h5_file_path):
            visualize_h5_event_representations(
                h5_file_path=h5_file_path,
                save_images=False,  # False=实时显示, True=保存图像
                vis_mode='stacked',  # 'stacked', 'channels', 'mixed_density', 'summary'
                max_frames=20,
                color_mode='viridis'
            )
        else:
            print(f"示例文件不存在: {h5_file_path}")